# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment2/question1/main.py
import matplotlib.pyplot
import numpy
import sympy
import sympy.abc


def _associated_laguerre(k, p_times_2, caches) -> sympy.Expr | int:
    key = (k, p_times_2)
    if key in caches:
        return caches[key]
    t = sympy.abc.t
    if k == 0:
        result = 1
    else:
        previous = _associated_laguerre(k - 1, p_times_2, caches)
        p = p_times_2 / 2
        result = sympy.diff(previous, t) * t
        result += (k + p - t) * previous
        result /= k
    caches[key] = result
    return result


def _draw_sympy_expression(expression, symbol):
    x = numpy.linspace(-100, 100, 1000)
    f = sympy.lambdify(symbol, expression)
    y = [f(item) for item in x]
    matplotlib.pyplot.plot(x, y)


def main():
    caches = {}
    plot_index = 1
    for n in range(0, 5 + 1):
        for l_value in range(n % 2, n + 1, 2):
            k = (n - l_value) // 2
            p_times_2 = l_value * 2 + 1
            result = _associated_laguerre(k, p_times_2, caches)
            result = sympy.expand(result)
            print(f"n={n},l={l_value}: {result}", end="    ")
            matplotlib.pyplot.subplots_adjust(hspace=0.5)
            matplotlib.pyplot.subplot(3, 4, plot_index)
            matplotlib.pyplot.title(rf"$L_{{{k}}}^{{{p_times_2}/2}}={{{sympy.latex(result)}}}$")
            _draw_sympy_expression(result, sympy.abc.t)
            plot_index += 1
        print()
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
